#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
NetPyNE 神经网络模型定义
用于 Geppetto 前端可视化
"""

from netpyne import specs, sim

class SimpleNetworkModel:
    """简单的神经网络模型"""
    
    def __init__(self):
        self.netParams = specs.NetParams()
        self.simConfig = specs.SimConfig()
        self.setup_network()
        self.setup_simulation()
    
    def setup_network(self):
        """设置网络参数"""
        
        # 网络参数
        self.netParams.sizeX = 100  # x 维度大小 (um)
        self.netParams.sizeY = 100  # y 维度大小 (um)
        self.netParams.sizeZ = 100  # z 维度大小 (um)
        
        # 细胞类型定义
        ## PYR 细胞 (锥体神经元)
        self.netParams.cellParams['PYR'] = {
            'secs': {
                'soma': {
                    'geom': {
                        'diam': 18.8,
                        'L': 18.8,
                        'Ra': 123.0
                    },
                    'mechs': {
                        'hh': {
                            'gnabar': 0.12,
                            'gkbar': 0.036,
                            'gl': 0.003,
                            'el': -70
                        }
                    }
                },
                'dend': {
                    'geom': {
                        'diam': 5.0,
                        'L': 150.0,
                        'Ra': 150.0
                    },
                    'topol': {
                        'parentSec': 'soma',
                        'parentX': 1.0,
                        'childX': 0
                    },
                    'mechs': {
                        'pas': {
                            'g': 0.0000357,
                            'e': -70
                        }
                    }
                }
            }
        }
        
        ## PV 细胞 (中间神经元)
        self.netParams.cellParams['PV'] = {
            'secs': {
                'soma': {
                    'geom': {
                        'diam': 15.0,
                        'L': 15.0,
                        'Ra': 120.0
                    },
                    'mechs': {
                        'hh': {
                            'gnabar': 0.15,
                            'gkbar': 0.04,
                            'gl': 0.003,
                            'el': -70
                        }
                    }
                }
            }
        }
        
        # 种群定义
        ## PYR 种群
        self.netParams.popParams['PYR_pop'] = {
            'cellType': 'PYR',
            'numCells': 20,
            'xRange': [0, 100],
            'yRange': [0, 100],
            'zRange': [0, 100]
        }
        
        ## PV 种群
        self.netParams.popParams['PV_pop'] = {
            'cellType': 'PV',
            'numCells': 5,
            'xRange': [0, 100],
            'yRange': [0, 100],
            'zRange': [0, 100]
        }
        
        # 突触机制
        self.netParams.synMechParams['AMPA'] = {
            'mod': 'Exp2Syn',
            'tau1': 0.1,
            'tau2': 5.0,
            'e': 0
        }
        
        self.netParams.synMechParams['GABA'] = {
            'mod': 'Exp2Syn',
            'tau1': 0.5,
            'tau2': 10.0,
            'e': -70
        }
        
        # 连接规则
        ## PYR -> PYR (兴奋性连接)
        self.netParams.connParams['PYR->PYR'] = {
            'preConds': {'pop': 'PYR_pop'},
            'postConds': {'pop': 'PYR_pop'},
            'probability': 0.1,
            'weight': 0.005,
            'delay': 5,
            'synMech': 'AMPA'
        }
        
        ## PYR -> PV (兴奋性连接)
        self.netParams.connParams['PYR->PV'] = {
            'preConds': {'pop': 'PYR_pop'},
            'postConds': {'pop': 'PV_pop'},
            'probability': 0.2,
            'weight': 0.01,
            'delay': 5,
            'synMech': 'AMPA'
        }
        
        ## PV -> PYR (抑制性连接)
        self.netParams.connParams['PV->PYR'] = {
            'preConds': {'pop': 'PV_pop'},
            'postConds': {'pop': 'PYR_pop'},
            'probability': 0.3,
            'weight': 0.01,
            'delay': 5,
            'synMech': 'GABA'
        }
        
        # 刺激
        self.netParams.stimSourceParams['bkg'] = {
            'type': 'NetStim',
            'rate': 10,
            'noise': 0.5
        }
        
        self.netParams.stimTargetParams['bkg->PYR'] = {
            'source': 'bkg',
            'conds': {'pop': 'PYR_pop'},
            'weight': 0.01,
            'delay': 5,
            'synMech': 'AMPA'
        }
    
    def setup_simulation(self):
        """设置仿真参数"""
        
        self.simConfig.duration = 1000  # 仿真时长 (ms)
        self.simConfig.dt = 0.025  # 时间步长 (ms)
        self.simConfig.verbose = False
        self.simConfig.recordTraces = {'V_soma': {'sec': 'soma', 'loc': 0.5, 'var': 'v'}}
        self.simConfig.recordStep = 0.1
        
        # 记录
        self.simConfig.recordCells = ['all']
        self.simConfig.recordStim = True
        self.simConfig.recordLFP = [[50, 50, 50]]
        
        # 分析
        self.simConfig.analysis['plotRaster'] = {
            'saveFig': False,
            'showFig': False
        }
        
        self.simConfig.analysis['plotTraces'] = {
            'include': [0, 1, 2],
            'saveFig': False,
            'showFig': False
        }
        
        self.simConfig.analysis['plot2Dnet'] = {
            'saveFig': False,
            'showFig': False
        }
    
    def create_network(self):
        """创建网络"""
        sim.initialize(self.netParams, self.simConfig)
        sim.net.createPops()
        sim.net.createCells()
        sim.net.connectCells()
        sim.net.addStims()
        return sim.net
    
    def run_simulation(self):
        """运行仿真"""
        sim.setupRecording()
        sim.runSim()
        return sim.allSimData
    
    def get_network_data(self):
        """获取网络数据（用于 Geppetto 可视化）"""
        if not hasattr(sim, 'net') or sim.net is None:
            self.create_network()
        
        network_data = {
            'cells': [],
            'connections': [],
            'populations': []
        }
        
        # 细胞数据
        for gid, cell in sim.net.cells.items():
            cell_data = {
                'gid': gid,
                'tags': cell.tags,
                'x': cell.tags.get('x', 0),
                'y': cell.tags.get('y', 0),
                'z': cell.tags.get('z', 0),
                'pop': cell.tags.get('pop', 'unknown'),
                'cellType': cell.tags.get('cellType', 'unknown')
            }
            network_data['cells'].append(cell_data)
        
        # 连接数据
        for cell in sim.net.cells.values():
            for conn in cell.conns:
                conn_data = {
                    'preGid': conn.get('preGid', -1),
                    'postGid': cell.gid,
                    'weight': conn.get('weight', 0),
                    'delay': conn.get('delay', 0),
                    'synMech': conn.get('synMech', 'unknown')
                }
                network_data['connections'].append(conn_data)
        
        # 种群数据
        for pop_name, pop in sim.net.pops.items():
            pop_data = {
                'name': pop_name,
                'cellType': pop.tags.get('cellType', 'unknown'),
                'numCells': len(pop.cellGids)
            }
            network_data['populations'].append(pop_data)
        
        return network_data
    
    def get_simulation_data(self):
        """获取仿真数据"""
        if not hasattr(sim, 'allSimData') or sim.allSimData is None:
            return None
        
        sim_data = {
            'spikes': [],
            'traces': {},
            'lfp': []
        }
        
        # Spike 数据
        if 'spkt' in sim.allSimData and 'spkid' in sim.allSimData:
            for t, gid in zip(sim.allSimData['spkt'], sim.allSimData['spkid']):
                sim_data['spikes'].append({'t': t, 'gid': gid})
        
        # Trace 数据
        if hasattr(sim.allSimData, 'keys'):
            for key in sim.allSimData.keys():
                if key.startswith('V_'):
                    sim_data['traces'][key] = sim.allSimData[key]
        
        # LFP 数据
        if 'LFP' in sim.allSimData:
            sim_data['lfp'] = sim.allSimData['LFP']
        
        return sim_data


def create_simple_model():
    """创建简单模型的便捷函数"""
    model = SimpleNetworkModel()
    return model


if __name__ == '__main__':
    # 测试模型
    print("创建 NetPyNE 模型...")
    model = create_simple_model()
    
    print("创建网络...")
    model.create_network()
    
    print("获取网络数据...")
    network_data = model.get_network_data()
    print(f"细胞数量: {len(network_data['cells'])}")
    print(f"连接数量: {len(network_data['connections'])}")
    print(f"种群数量: {len(network_data['populations'])}")
    
    print("\n运行仿真...")
    sim_data = model.run_simulation()
    
    print("获取仿真数据...")
    simulation_data = model.get_simulation_data()
    print(f"Spike 数量: {len(simulation_data['spikes'])}")
    
    print("\n模型测试完成！")

